import asyncio
from typing import List

from llama_index.core.agent.workflow import  FunctionAgent
from llama_index.core.base.llms.types import ChatMessage
from llama_index.core.indices.vector_store import VectorIndexRetriever
from llama_index.core.storage.chat_store.sql import SQLAlchemyChatStore
from llama_index.core.tools import QueryEngineTool
from llama_index.core.vector_stores import SimpleVectorStore
from llama_index.core.schema import  TextNode
from llama_index.core import Settings, SimpleKeywordTableIndex, SummaryIndex, get_response_synthesizer, \
    VectorStoreIndex, StorageContext
from llama_index.embeddings.zhipuai import ZhipuAIEmbedding
from llama_index.core.graph_stores import SimplePropertyGraphStore
from llama_index.core.schema import Document
from pydantic import BaseModel



embed_model = ZhipuAIEmbedding(
    model="embedding-2",
    api_key="f387f5e4837d4e4bba6d267682a957c9.PmPiTw8qVlsI2Oi5"
    # With the `embedding-3` class
    # of models, you can specify the size
    # of the embeddings you want returned.
    # dimensions=1024
)
Settings.embed_model=embed_model

from llama_index.llms.deepseek import DeepSeek

llm = DeepSeek(model="deepseek-chat", api_key="sk-605e60a1301040759a821b6b677556fb")
Settings.llm = llm


nodes = [TextNode(text='''‌法律条款检索‌：合并分散的关联条款（如"合同法第12条"与"司法解释第3条"），输出完整法律依据''')]
nodes001 = [TextNode(text='''‌医疗文献分析‌：当查询"糖尿病治疗方案"时，自动合并药物说明、临床实验数据等关联段落''')]

vectorStoreIndex=VectorStoreIndex(nodes=nodes)

base_retriever = VectorIndexRetriever(vectorStoreIndex)

vectorStoreIndex001=VectorStoreIndex(nodes=nodes001)

base_retriever001 = VectorIndexRetriever(vectorStoreIndex001)

from llama_index.core.retrievers import QueryFusionRetriever

vector_retriever = vectorStoreIndex.as_retriever(similarity_top_k=5)
bm25_retriever = vectorStoreIndex001.as_retriever( similarity_top_k=5)

fusion_retriever = QueryFusionRetriever(
    retrievers=[vector_retriever, bm25_retriever],
    llm=llm,
    mode='reciprocal_rerank',
    similarity_top_k=2,
)
results = fusion_retriever.retrieve("条款和‌医疗文献分析‌")
print(results)
print(len(results))